import os
#from CalTorque import *
from datetime import datetime
from string import Template
import numpy
import matplotlib
matplotlib.use('agg')
from matplotlib import pyplot as PLT
import cPickle as pickle
from collections import Counter
from decimal import Decimal
import datetime

def SaveAveragePKL(x, y):
    avgdata = {}
    avgdata['x'] = x
    avgdata['y'] = y
    output = open('pickled/pickled_average_in.pkl', 'wb')
    pickle.dump(avgdata, output, -1)
    output.close()

def AverageIN(database, bins, max_d, saveAvg = False):

    # get today's date for file title
    d = datetime.datetime.now()
    q = str(d.strftime('%Y%m%d'))
    
    
    # bin and average data from database
    x_inputs = numpy.zeros(bins)
    X = numpy.linspace(0,max_d,bins)
    Y = numpy.zeros(bins)
    Y_std = numpy.zeros(bins)
    Y2 = []
    
    for i in range(X.size):
        Y2.append([])
        
    for run in database:
       if run.Direction() != 'in': continue #Skip all retractions
       if run.Source() != 'weakTh-228': continue #Skip all other sources
       if run.Station() in ['S8', 'S11', 'Other']: continue # Skip calibration runs
       print run
       position = run.NewPosition()
       torque = run.Torque()
       
       # bin
       for j,x in enumerate(X):
          pos_torque = []
          x0 = x-(max_d/bins)/2.
          x1 = x+(max_d/bins)/2.
          for i, p in enumerate(position):
             if p > x1: continue  #TODO test this
             if p>x0 and p<=x1:
                pos_torque.append(torque[i])
          if len(pos_torque)>0:
             if max(pos_torque) == 0: continue
             Y[j] += max(pos_torque)
             Y2[j].append(max(pos_torque))
             x_inputs[j]+=1.
    
    # calculate average for each bin
    for i,n in enumerate(x_inputs):
       if n==0: continue
       Y[i] = Y[i]/n
       Y_std[i] = numpy.std(Y2[i], dtype = numpy.float64)
    
    if saveAvg:
        SaveAveragePKL(X, Y)
    
    # plot average
    fig = PLT.figure(figsize=(15,8),dpi=150)
    box = [0.14, 0.14, 0.76, 0.76]
    ax1 = fig.add_axes(box)
    ax1.set_xlabel('Position (cm)')
    ax1.set_ylabel('Torque (in.oz)')
    miny = min(Y_std)
    maxy = max(torque)
    maxx = max(position)
    
    ## Add station locations
    if maxx >= 328.9:
        ax1.axvline(338.9, color = 'r', linewidth = 1, linestyle = '-')
        ax1.annotate('S2', xy = (340.9, miny * 0.8), color = 'r')
    if maxx >= 382.6:
        ax1.axvline(392.6, color = 'r', linewidth = 1, linestyle = '-')
        ax1.annotate('S5', xy = (394.6, miny * 0.8), color = 'r')
    if maxx >= 435.7:
        ax1.axvline(445.7, color = 'r', linewidth = 1, linestyle = '-')
        ax1.annotate('S8', xy = (447.7, miny * 0.8), color = 'r')
    if maxx >= 491.7:
        ax1.axvline(501.7, color = 'r', linewidth = 1, linestyle = '-')
        ax1.annotate('S11', xy = (503.7, miny * 0.8), color = 'r')
    if maxx >= 595.8:
        ax1.axvline(605.8, color = 'r', linewidth = 1, linestyle = '-')
        ax1.annotate('S17', xy = (607.8, miny * 0.8), color = 'r')
    
    ## add bend locations
    if maxx >= 47.8:
        ax1.axvspan(47.8, 93.5, color = 'g', alpha = 0.2)
        ax1.annotate('Bellows', xy = (48, miny), color = 'g')
    if maxx >= 136.4:
        ax1.axvspan(136.4, 144.6, color = 'g', alpha = 0.2)
        ax1.annotate('First outer bend', xy = (136, miny), color = 'g')
    if maxx >= 196.5:
        ax1.axvspan(196.5, 218.0, color = 'g', alpha = 0.2)
        ax1.annotate('The "S"', xy = (212, miny), color = 'g')
    if maxx >= 228.3:
        ax1.axvspan(228.3, 241.5, color = 'g', alpha = 0.2)
    if maxx >= 274.6:
        ax1.axvspan(274.6, 284.6, color = 'g', alpha = 0.2)
        ax1.annotate('Bend A5', xy = (274, miny), color = 'g')
    if maxx >= 309.0:
        ax1.axvspan(309.0, 318.9, color = 'g', alpha = 0.2)
        ax1.annotate('Bend A6', xy = (309, miny), color = 'g')
    if maxx >= 358.8:
        ax1.axvspan(358.8, 370.8, color = 'g', alpha = 0.2)
        ax1.annotate('Bend A7', xy = (359, miny), color = 'g')
    if maxx >= 414.5:
        ax1.axvspan(414.5, 426.5, color = 'g', alpha = 0.2)
        ax1.annotate('Bend A8', xy = (414, miny), color = 'g')
    if maxx >= 445.3:
        ax1.axvspan(445.3, 455.3, color = 'g', alpha = 0.2)
        ax1.annotate('Bend A9', xy = (445, miny), color = 'g')
    if maxx >= 467.5:
        ax1.axvspan(467.5, 479.5, color = 'g', alpha = 0.2)
        ax1.annotate('Bend A10', xy = (467, miny), color = 'g')
    if maxx >= 523.8:
        ax1.axvspan(523.8, 535.8, color = 'g', alpha = 0.2)
        ax1.annotate('Bend A11', xy = (523, miny), color = 'g')
    if maxx >= 571.6:
        ax1.axvspan(571.6, 583.6, color = 'g', alpha = 0.2)
        ax1.annotate('Bend A12', xy = (571, miny), color = 'g')
    if maxx >= 627.9:
        ax1.axvspan(627.9, 639.9, color = 'g', alpha = 0.2)
        ax1.annotate('Bend A13', xy = (627, miny), color = 'g')
    if maxx >= 652.7:
        ax1.axvspan(652.7, 660.7, color = 'g', alpha = 0.2)
        ax1.annotate('Bend A14', xy = (652, miny), color = 'g')
    if maxx >= 678.5:
        ax1.axvspan(678.5, 681.1, color = 'g', alpha = 0.2)
        ax1.annotate('Bend A15', xy = (678, miny), color = 'g')
    
    PLT.title('Average of all IN deployments')
    ax1.grid()
    PLT.errorbar(X, Y, xerr = None, yerr = Y_std)
    fig.savefig('averagePlots/average_in_model_Update-' + q + '.png')